package com.star.spring.ai.controller;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.time.LocalDateTime;

@RestController
@RequestMapping("/chat")
public class ChatClientController {

    @Autowired
    private ChatClient chatClient;

    @GetMapping("/string")
    String generation(String userInput) {
        return this.chatClient.prompt().user(userInput).call().content();
    }

    @GetMapping(value = "/stream", produces = "text/html;charset=UTF-8")
    Flux<String> streamGeneration(String userInput) {
        Flux<String> content = chatClient
                .prompt()
                .user(userInput)
                .stream()
                .content();

        // 中断长连接,点断收到interrupted就中断长连接
        return content.concatWith(Flux.just("interrupted"));
    }

    @GetMapping(value = "/role", produces = "text/html;charset=UTF-8")
    Flux<String> withParam(String userInput) {
        return chatClient
                .prompt()
                // 传入当天日期
                .system(promptSystemSpec ->
                        promptSystemSpec.param("current_date", LocalDateTime.now().toString()))

                .user(userInput).stream().content();
    }

    @GetMapping(value = "/memory", produces = "text/html;charset=UTF-8")
    Flux<String> withMemory(String userInput) {
        return chatClient
                .prompt()
                // 传入当前时间
                .system(promptSystemSpec ->
                        promptSystemSpec.param("current_date", LocalDateTime.now().minusDays(1).toString()))

                // 传入记忆最大数量
                .advisors(advisorSpec -> advisorSpec.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                .functions("functionChange")

                .user(userInput).stream().content();
    }


}
